import cv2
import numpy as np
import subprocess
import rospy
from sensor_msgs.msg import Image
from cv_bridge import CvBridge, CvBridgeError

# FFmpeg 推流命令
# 替换 <server-ip> 和 <stream-key> 为你的 RTMP 服务器地址和流密钥
ffmpeg_command = [
    'ffmpeg',
    '-y',
    '-f', 'rawvideo',
    '-vcodec', 'rawvideo',
    '-pix_fmt', 'bgr24',
    '-s', '640x480',  # 尺寸
    '-r', '100',  # 帧率
    '-i', '-',  # 标准输入
    '-c:v', 'libx264',
    '-preset', 'veryfast',
    '-tune', 'zerolatency',
    '-f', 'flv',
    'rtmp://124.221.6.242:1935/live/rtmp'  # 替换为你的 RTMP 服务器地址
]

class ROS2FFMPEGStreamer:
    def __init__(self):
        self.bridge = CvBridge()
        self.process = subprocess.Popen(ffmpeg_command, stdin=subprocess.PIPE)
        rospy.init_node('ros2ffmpeg_streamer', anonymous=True)
        rospy.Subscriber('/usb_cam/image_raw', Image, self.image_callback)
        rospy.spin()

    def image_callback(self, msg):
        try:
            # 将 ROS 图像消息转换为 OpenCV 图像
            cv_image = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
            # 将图像写入 FFmpeg 的 stdin
            self.process.stdin.write(cv_image.tobytes())
        except CvBridgeError as e:
            rospy.logerr(f"CV Bridge Error: {e}")

    def __del__(self):
        self.process.stdin.close()
        self.process.wait()

if __name__ == '__main__':
    try:
        ROS2FFMPEGStreamer()
    except rospy.ROSInterruptException:
        pass
